'''
6.代理（Agent） 定义：具备自主决策能力的AI实体，能感知环境、规划、调用工具、执行动作并迭代。

应用场景：
AI助手：如Auto-GPT、BabyAGI，能自主完成“调研竞品 → 写报告 → 发邮件”等任务。
自动化客服代理：理解用户问题 → 查询知识库 → 调用API → 生成回复。
智能机器人：结合视觉、语音、动作，完成物理世界任务（如仓储机器人）。
金融交易代理：监控市场 → 分析数据 → 决策买卖 → 执行交易。
游戏AI：NPC根据玩家行为动态调整策略。
✅ 优势：高度自主、适应复杂环境、支持长期目标达成。
''' 

# 导入必要的库和模块
from langchain_core.tools import tool
from langgraph.graph import MessagesState
from langchain_core.messages import SystemMessage, HumanMessage, ToolMessage
from langgraph.graph import StateGraph, START, END
from typing_extensions import Literal
from utils import init_llm, set_env, visualize_graph

@tool
def multiply(a: int, b: int) -> int:
    """Multiply a and b.

    Args:
        a: first int
        b: second int
    """
    return a * b


@tool
def add(a: int, b: int) -> int:
    """Adds a and b.

    Args:
        a: first int
        b: second int
    """
    return a + b


@tool
def divide(a: int, b: int) -> float:
    """Divide a and b.

    Args:
        a: first int
        b: second int
    """
    return a / b


def llm_call(state: MessagesState):
    """LLM 决定是否调用工具"""

    return {
        "messages": [
            llm_with_tools.invoke(
                [
                    SystemMessage(
                        content="您是一位能力强大的助手，负责对一组输入进行算术运算。"
                    )
                ]
                + state["messages"]
            )
        ]
    }


def tool_node(state: dict):
    result = []
    for tool_call in state["messages"][-1].tool_calls:
        tool = tools_by_name[tool_call["name"]]
        observation = tool.invoke(tool_call["args"])
        result.append(ToolMessage(content=observation, tool_call_id=tool_call["id"]))
    return {"messages": result}


def should_continue(state: MessagesState) -> Literal["environment", END]:
    messages = state["messages"]
    last_message = messages[-1]
    if last_message.tool_calls:
        return "Action"
    return END

    
def build_graph():
    """
    构建工作流图函数
    
    功能：
    1. 创建状态图实例
    2. 添加所有节点（编排器、工作节点、综合器）
    3. 定义节点之间的连接关系
    4. 编译并返回可执行的工作流图
    
    Returns:
        编译后的工作流图
    """
    
    # 创建状态图实例
    agent_builder = StateGraph(MessagesState)

    # 添加各个节点
    agent_builder.add_node("llm_call", llm_call)
    agent_builder.add_node("environment", tool_node)


    agent_builder.add_edge(START, "llm_call")
    agent_builder.add_conditional_edges(
        "llm_call",
        should_continue,
        {
            "Action": "environment",
            END: END,
        },
    )
    agent_builder.add_edge("environment", "llm_call")

    agent = agent_builder.compile()

    return agent

if __name__ == "__main__":
    # 主程序入口
    set_env()  # 设置环境变量
    llm = init_llm(temperature=0)  # 初始化LLM，温度设为0以获得确定性输出
    tools = [add, multiply, divide]
    tools_by_name = {tool.name: tool for tool in tools}
    llm_with_tools = llm.bind_tools(tools)
    graph = build_graph()  # 构建工作流图
 
    # visualize_graph(graph)  # 可视化工作流图（已注释）

    messages = [HumanMessage(content="3 加 4 等于几？")]
    messages = graph.invoke({"messages": messages})
    for m in messages["messages"]:
        m.pretty_print()






